package com.agentai.model.controller;


import com.agentai.framework.common.utils.Result;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.Media;
import org.springframework.ai.model.ModelRequest;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.util.MimeType;
import org.springframework.util.MimeTypeUtils;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
import java.net.URL;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;

@Slf4j
@RequestMapping("/ollama")
@RestController
public class ChatDemoController {

    @GetMapping("/deep/seek/message")
    public Result<ChatResponse> message(@RequestParam("message") String message) {
        message = "请使用中文简体回答：" + message;
        Prompt prompt = new Prompt(new UserMessage(message));
        ChatResponse response = ollamaChatModel.call(prompt);
        return Result.ok(response);
    }

    @GetMapping("/deep/seek/stream")
    public Flux<ChatResponse> stream(@RequestParam("message") String message) {
        message = "请使用中文简体回答：" + message;
        Prompt prompt = new Prompt(new UserMessage(message));
        OllamaOptions options = new OllamaOptions();
        options.setModel("deepseek-r1:7b");
        return print(options, prompt);
    }

    @GetMapping("/img/stream")
    public Flux<ChatResponse> imgStream(@RequestParam("message") String message) throws Exception {
//        message = "请使用中文简体回答：" + message;
        URL url = new URL("http://localhost:8080/upload/20250315/ocr_50819.jpg");

        String base64Image = convert("C:\\Users\\lwh\\Desktop\\pdf\\ocr.png");

        List<Media> medias = new ArrayList<>();

        Media media = Media.builder().mimeType(MimeTypeUtils.IMAGE_PNG).data(base64Image).build();

        medias.add(media);

        UserMessage userMessage = new UserMessage(message,medias);

        Prompt prompt = new Prompt(userMessage);
        OllamaOptions options = new OllamaOptions();
        options.setModel("llama3.2-vision:11b");
        return print(options, prompt);
    }

    @Resource
    private OllamaChatModel ollamaChatModel;

    private Flux<ChatResponse> print(OllamaOptions options, Prompt prompt) {

        OllamaApi api = new OllamaApi();
        OllamaChatModel test = OllamaChatModel.builder().ollamaApi(api).defaultOptions(options).build();

        Flux<ChatResponse> chatResponseFlux = test.stream(prompt);
        chatResponseFlux
                .doOnNext(response -> {
                    for (Generation result : response.getResults()) {
                        log.info("Received response: {}", result.getOutput().getText());
                    }
                })
                .subscribe();

        return chatResponseFlux;
    }

    public static String convert(String imagePath) throws Exception {
        // 1. 读取图片文件
        BufferedImage image = ImageIO.read(new File(imagePath));

        // 2. 创建字节输出流
        ByteArrayOutputStream baos = new ByteArrayOutputStream();

        // 3. 获取文件扩展名（自动识别格式）
        String formatName = imagePath.substring(imagePath.lastIndexOf(".") + 1).toLowerCase();

        // 4. 将图像写入字节流
        if (!ImageIO.write(image, formatName, baos)) {
            throw new IllegalArgumentException("不支持的图片格式: " + formatName);
        }

        // 5. 转换为Base64字符串
        return Base64.getEncoder().encodeToString(baos.toByteArray());
    }

}
